from psychsim.pwl import *
from psychsim.action import Action,ActionSet
from psychsim.world import World,stateKey
from psychsim.agent import Agent
from psychsim.reward import *

''' 
This example is used to test the update of rewards.
'''

class UpdateRewardsTestCase(unittest.TestCase):

    def setUp(self):
        #init world and agent, set maxRounds and turn order
        self.world = World()
        self.claudius = Agent('Claudius')
        self.world.addAgent(self.claudius)

        self.maxRounds=1
        self.world.setOrder(['Claudius'])

    def addStates(self):    
        #add states for round and poisoned                
        self.world.defineState(None,'round',int)
        self.world.setState(None,'round',0)
        
        self.world.defineState(None,'poisoned',bool)
        self.world.setState(None,'poisoned',False)
                            
    def addTermination(self):
        #add termination if at max round
        self.world.addTermination(makeTree({'if': thresholdRow(stateKey(None,'round'),self.maxRounds),
                                   True: True, False: False}))

    def addActions(self):
        #add actions and dynamics
        action = self.claudius.addAction({'verb': 'poisons'})
        tree = makeTree(setTrueMatrix(stateKey(None,'poisoned')))
        self.world.setDynamics('poisoned',action,tree)
        tree = makeTree(incrementMatrix(stateKey(None,'round'),1))
        self.world.setDynamics('round',action,tree)     
        action =  self.claudius.addAction({'verb': 'do nothing'})
        self.world.setDynamics('round',action,tree) 

    def addRewards(self):
        #set rewards for claudius
        tree = maximizeFeature(stateKey(None,'poisoned'))
        self.claudius.setReward(tree, 1.0, True)

    def specifyModel(self):
        #make sure to update model
        self.world.setModel('Claudius', True)
                                                       
    def testSetup(self):
        self.addStates()
        self.addTermination()
        self.addActions()
        self.addRewards()
        self.specifyModel()

    def testDoubleSetRewards(self):
        #check setting the same reward twice
        tree = maximizeFeature(stateKey(None,'poisoned'))
        self.claudius.setReward(tree, 1.0, True)
        reward1 = self.claudius.getAttribute('R', True)
        self.claudius.setReward(tree, 1.0, True)
        reward2 = self.claudius.getAttribute('R', True)
        self.assertEqual(reward1, reward2)

        #check setting a new reward
        self.claudius.setReward(tree,2.0,True)
        reward3 = self.claudius.getAttribute('R', True)
        self.assertEqual(str(reward1.keys()[0]), str(reward3.keys()[0]))
        self.assertEqual(reward3.values()[0], 2.0)
            
    def testRun(self):
        self.testSetup()
        self.world.printState(beliefs=True)
        #self.world.printModels()
        for t in range(self.maxRounds + 1):
            self.world.explain(self.world.step(),level=4)
            self.world.printState(beliefs=True)
            #self.world.printModels()
            if self.world.terminated():
                break 
            
if __name__ == "__main__":
    suite = unittest.TestSuite()
    for method in dir(UpdateRewardsTestCase):
       if method.startswith("test"):
          suite.addTest(UpdateRewardsTestCase(method))
    unittest.TextTestRunner().run(suite)